#!/usr/bin/python
import astro.spec
import matplotlib.pyplot as pl
from astro.utilities import readatom, findtrans
from astro.io_util import readheader
from astro.io import readtxt
import plotvel
import numpy as np

#import numpy as np

# read in one or more spectra.  left and right arrows to switch between them

atom = readatom('/home/nhmc/installed/vpfit/atom.dat')

usage = """\
Usage:  plotspec.py filename1 filename2 ...

A file lines.txt in the working directory can be used to give a list
of lines and rest wavelengths.  e.g.

# name       wavelength
H-a          6564.0

and so on ...
"""

help = """Left and right arrows to switch between spectra.
"""

REST = dict(LyLim   = 912.0    ,
            CIIIa   = 977.020  ,
            Lyb     = 1025.72  ,
            Lya     = 1215.6701,
            NV      = 1240.14  ,
            OI_SiII = 1303.0   ,
            CII     = 1334.5323,
            SiIV    = 1400.    ,
            CIV     = 1550.    ,
            CIIIb   = 1908.73  ,
            MgII    = 2798.75  ,
            OII     = 3728.48 ,
            Hg      = 4341.7   ,
            Hb      = 4862.    ,
            Ha      = 6564.61  ,
            )
REST = np.rec.fromrecords(REST.items(), names='ion,wa')
REST.sort(order='wa')

class Junk(object):
    def __init__(self, filenames, fig, lines=None):
        self.filenames = []
        self.filenames.extend(filenames)
        self.i = 0
        self.zem_p1 = None
        self.n = len(self.filenames)
        self.spec = [None] * self.n          # cache spectra
        self.info = [None] * self.n          # and spec info
        if lines is None:
            lines = REST
        self.lines = lines
        temp = range(1, len(lines)+1)
        self.linehelp = '\n'.join('%i: %s %.4f' % (i,r['ion'],r['wa']) for
                                  i,r in zip(temp, lines))
        self.linehelp += '\nEnter transition: '
        self.cids = dict()
        # disable any existing key press callbacks
        cids = list(fig.canvas.callbacks.callbacks['key_press_event'])
        for cid in cids:
            fig.canvas.callbacks.disconnect(cid)
        self.connect(fig)
        self.fig = fig

    def update(self):
        i = self.i
        filename = self.filenames[i]
        if self.spec[i] is None:
            print 'Reading %s' % filename
            self.info[i] = readheader(filename)
            self.spec[i] = astro.spec.read(filename)
        self.artists = []
        self.fig.clf()
        a = self.fig.add_subplot(111)
        self.spec[i].plot(ax=a, show=0)
        a.set_title(filename)
        a.set_autoscale_on(False)
        if set(self.info[i]).intersection(['z', 'mag']):
            if 'note' in self.info[i]:
                text = 'z: %(z)s mag: %(mag)s\nnote:%(note)s' % self.info[i]
            else:
                text = 'z: %(z)s mag: %(mag)s' % self.info[i]
            a.text(0.5, 0.85, text, transform=a.transAxes)
            self.plotlines(float(self.info[i]['z'])+1)
        astro.spec.plotatmos(a)
        pl.draw()

    def plotlines(self, zp1):
        if self.artists:
            for l in self.artists:
                try:
                    l.remove()
                except ValueError:
                    # plot has been removed
                    pass
        self.artists = astro.spec.plotlines(zp1-1, pl.gca(), lines=self.lines)
        pl.draw()

    def on_keypress_navigate(self, event):
        if event.key == 'right':
            if self.i == self.n-1:
                print 'At the last spectrum.'
                return
            self.i += 1
            self.update()
        elif event.key == 'left':
            if self.i == 0:
                print 'At the first spectrum.'
                return
            self.i -= 1
            self.update()
    
    def on_keypress_plotz(self, event):
        ax = event.inaxes
        if ax is None:  return
        if event.key == 'i':
            # id line to get redshift
            while True:
                c = raw_input(self.linehelp)
                try:
                    i = int(c) - 1
                    ion = self.lines[i]['ion']
                    wa = self.lines[i]['wa']
                except (TypeError, IndexError):
                    continue
                else:
                    break

            zp1 = event.xdata / wa
            print 'z=%.3f, %s %.2f' % (zp1-1, ion, wa)
        else:
            return
        self.zem_p1 = zp1
        self.plotlines(zp1)
        
    def on_keypress_plotvel(self, event):
        if event.key == 'v':
            if self.zem_p1 is None:
                print 'You must assign a redshift'
                return
            spec = self.spec[self.i]
            z = self.zem_p1 - 1.
            fig = pl.figure(figsize=(14,11))
            junk = plotvel.VelplotWrap(spec.wa, spec.fl/spec.co,
                                       spec.er/spec.co,
                                       fig, z, spec.filename)
            print plotvel.help

    def connect(self, fig):
        cids = dict()
        cids['nav'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_navigate)
        cids['plotz'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_plotz)
        cids['plotv'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_plotvel)
        self.cids.update(cids)

    def diconnect(self):
        for key in ['nav', 'plotz']:
            self.fig.canvas.mpl_disconnect(self.cids[key])

def main(args):
    if len(args) < 1:
        print usage
        return
    try:
        fh = open('lines.txt')
    except IOError:
        lines = None
    else:
        lines = readtxt(fh, names='ion,wa')

    fig = pl.figure(figsize=(13,5))
    ax = pl.axes([0.06,0.12,0.90,0.75])
    print help
    junk = Junk(args, fig, lines=lines)
    junk.update()
    pl.show()

if __name__ == '__main__':
    import sys
    sys.exit(main(sys.argv[1:]))
